import math
import random

import numpy as np
from spatialmath import SE3, SO3

from src.geometry import Ellipsoid

from .node import Node
from .rrt_map import RRTMap
from .rrt_parameter import RRTParameter
from .rrt_star_planner import RRTStarPlanner


class InformedRRTStarPlanner(RRTStarPlanner):
    def __init__(self, rrt_map: RRTMap, rrt_parameter: RRTParameter) -> None:
        self.c_min = np.linalg.norm(rrt_parameter.goal - rrt_parameter.start)
        self.x_center = (rrt_parameter.goal + rrt_parameter.start) / 2.0

        id1_t = np.array([1.0, 0.0, 0.0]).reshape(1, 3)
        a1 = (rrt_parameter.goal - rrt_parameter.start).reshape(3, 1) / self.c_min
        M = a1 @ id1_t
        U, S, Vh = np.linalg.svd(M, True, True)
        self.C = np.dot(np.dot(U, np.diag([1.0, 1.0, np.linalg.det(U) * np.linalg.det(np.transpose(Vh))])), Vh)
        self.r = np.array([0.0, 0.0, 0.0])

        super().__init__(rrt_map, rrt_parameter)

    def sample(self) -> Node:
        if not self.success:
            return super().sample()
        while True:
            self.r[0] = self.last_path_length / 2.0
            self.r[1] = np.sqrt(self.last_path_length ** 2 - self.c_min ** 2) / 2.0
            self.r[2] = self.r[1]
            L = np.diag(self.r)
            x_ball = InformedRRTStarPlanner.sample_unit_ball()
            rnd = np.squeeze(np.dot(np.dot(self.C, L), x_ball)) + self.x_center
            if self.in_area(rnd):
                return Node(rnd)

    @staticmethod
    def sample_unit_ball() -> np.ndarray:

        radius = random.random()
        theta = random.random()
        phi = random.random()

        return np.array([[np.cos(theta * 2 * np.pi) * np.sin(phi * np.pi)],
                         [np.sin(theta * 2 * np.pi) * np.sin(phi * np.pi)],
                         [np.cos(phi * np.pi)]]) * radius

    def draw_others(self, ax) -> None:
        if not self.success:
            return
        ellipsoid = Ellipsoid(SE3(*self.x_center) * SE3(SO3(self.C)), self.r)
        ellipsoid.plot(ax, c='y')
